""" 创建数据库连接
@Project: Rebar
@File: init.py
@Author: hzc
@Date: 2023/3/22 11:07
"""
import time
from contextlib import contextmanager

from sqlalchemy.ext.declarative import declarative_base

from config import MYSQL_CONFIG
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session

from model.db.well_model import Base as wellBase
from model.db.user_model import Base as userBase
from model.db.sys_model import Base as sysBase
from model.db.experiment_model import Base as experimentBase


def create_tables(engine_):
    wellBase.metadata.create_all(engine_)
    userBase.metadata.create_all(engine_)
    sysBase.metadata.create_all(engine_)
    experimentBase.metadata.create_all(engine_)


def create_project_tables(engine_):
    wellBase.metadata.create_all(engine_)
    experimentBase.metadata.create_all(engine_)


# 公共库数据库访问地址
SQLALCHEMY_DATABASE_URL = "mysql+pymysql://{USERNAME}:{PASSWORD}@{HOST}:{PORT}/{DATABASE}".format(**MYSQL_CONFIG)
# 创建数据库引擎
engine = create_engine(
    SQLALCHEMY_DATABASE_URL,
    pool_recycle=3600,
    pool_size=100,
)

# 创建数据库会话
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
# 创建公共session对象:
reservoir_sensitivity = SessionLocal()
# 创建表
create_tables(engine)
# 私有项目session对象（中间件中赋值） 默认公共数据库
projSession = reservoir_sensitivity


# 创建特定的数据库引擎池
def create_tenant_engine(db_name: str):
    tenant_db_url = "mysql+pymysql://{USERNAME}:{PASSWORD}@{HOST}:{PORT}/".format(**MYSQL_CONFIG) + db_name
    return create_engine(tenant_db_url, pool_recycle=3600,
                         pool_size=100, )


# 在每个请求中，选择正确的租户特定引擎和会话
def tenant_session_scope(db_name: str):
    tenant_engine = create_tenant_engine(db_name)
    SessionLocalTenant = sessionmaker(autocommit=False, autoflush=False, bind=tenant_engine)
    scoped_session_maker = scoped_session(SessionLocalTenant)
    session = scoped_session_maker()
    try:
        session.commit()
        return session
    except Exception as e:
        session.rollback()
        session.close()
        raise e
